{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\rwr\\.conda\\envs\\pytorch\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "# 加载指标计算方法\n",
    "from utils import *\n",
    "import pyiqa\n",
    "from FFT import *\n",
    "if torch.cuda.is_available():\n",
    "    device = 'cuda'\n",
    "else:\n",
    "    device = 'cpu'\n",
    "print('Using device:', device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建pyiqa指标对象\n",
    "brisque_metric = pyiqa.create_metric('brisque',device=device)\n",
    "ilniqe_metric = pyiqa.create_metric('ilniqe',device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置数据集的路径和下载选项\n",
    "Dataset_dir = '../data'\n",
    "Output_fil_dir = './results_filtered'\n",
    "train_metrics_dir = os.path.join(Output_fil_dir, 'train_metrics_filtered.csv')\n",
    "test_metrics_dir = os.path.join(Output_fil_dir, 'test_metrics_filtered.csv')\n",
    "#如果输出目录不存在，则创建\n",
    "if not os.path.exists(Output_fil_dir):\n",
    "    os.makedirs(Output_fil_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "#定义数据预处理\n",
    "transform = transforms.Compose([transforms.ToTensor()]) \n",
    "# 读取CITAR-100数据集\n",
    "trainset = torchvision.datasets.CIFAR100(root=Dataset_dir, train=True, download=True, transform=transform)\n",
    "testset = torchvision.datasets.CIFAR100(root=Dataset_dir, train=False, download=True, transform=transform)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 归一化滤波器后的array\n",
    "def normalize(gray_img_low, gray_img_high):\n",
    "    # gray_img_low, gray_img_high = filter_image(gray_image_img, filter_size=2, show_filters=False)\n",
    "    low_min, low_max = np.min(gray_img_low), np.max(gray_img_low)\n",
    "    high_min, high_max = np.min(gray_img_high), np.max(gray_img_high)\n",
    "    gray_img_low = (gray_img_low - low_min) / (low_max - low_min)\n",
    "    gray_img_high = (gray_img_high - high_min) / (high_max - high_min)\n",
    "    return gray_img_low, gray_img_high"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_img_filtered_metrics(cifar100_dataset, Output_fil_dir):\n",
    "    print(\"Calculating image metrics...\")\n",
    "    # 定义将张量转换为PIL图像的函数\n",
    "    to_pil = transforms.ToPILImage()\n",
    "    results_list = []  # 使用列表来收集数据\n",
    "\n",
    "    # 使用tqdm来显示进度\n",
    "    for i in tqdm(range(len(cifar100_dataset))):\n",
    "        \n",
    "        # 读取图像数据(张量形式)\n",
    "        image_ten, label = cifar100_dataset[i]\n",
    "        # 将张量转为PIL图像\n",
    "        image_img = to_pil(image_ten)\n",
    "        gray_img = image_img.convert('L')  # 将PIL图像转为灰度图像\n",
    "        # channels = image_ten.shape[0]\n",
    "        # print(channels)\n",
    "        gray_img_low, gray_img_high = filter_image(gray_img, filter_size=2, show_filters=False)\n",
    "        # # 将PIL图像转为灰度图像\n",
    "        # gray_img_low = low_img.convert('L') \n",
    "        # gray_img_high = high_img.convert('L') \n",
    "\n",
    "        # 将灰度图像转为numpy数组\n",
    "        img_int = np.array(gray_img).astype(np.int32)\n",
    "        img_int_low = np.array(gray_img_low).astype(np.int32)\n",
    "        img_double_low = np.array(gray_img_low).astype(np.float32)\n",
    "        img_double = np.array(gray_img).astype(np.float32)\n",
    "        img_int_high = np.array(gray_img_high).astype(np.int32)\n",
    "        img_double_high = np.array(gray_img_high).astype(np.float32)\n",
    "\n",
    "        img_ten_low = torch.from_numpy(gray_img_low)\n",
    "        img_ten_high = torch.from_numpy(gray_img_high)\n",
    "        # # 计算pyiqa指标\n",
    "        # brisque_score = brisque_metric(image_tensor.to(device)).item()\n",
    "        #ilniqe_score = ilniqe_metric(image_tensor.to(device)).item()\n",
    "        gray_img_low_norm, gray_img_high_norm = normalize(gray_img_low, gray_img_high)\n",
    "        image_tensor_low = torch.from_numpy(gray_img_low_norm)\n",
    "        image_tensor_high = torch.from_numpy(gray_img_high_norm)\n",
    "        # 将图像张量扩展为四维张量，单通道\n",
    "        image_tensor=image_ten.unsqueeze(0)\n",
    "        image_tensor_low_1ch = image_tensor_low.unsqueeze(0).unsqueeze(0)\n",
    "        image_tensor_high_1ch = image_tensor_high.unsqueeze(0).unsqueeze(0)\n",
    "        # 计算brisque指标\n",
    "        brisque_score = brisque_metric(image_tensor.to(device)).item()\n",
    "        brisque_score_low = brisque_metric(image_tensor_low_1ch.to(device)).item()\n",
    "        brisque_score_high = brisque_metric(image_tensor_high_1ch.to(device)).item()\n",
    "        # # 计算ilniqe指标，需要三通道图像\n",
    "        # image_tensor_low_3ch = image_tensor_low.unsqueeze(0).expand(3,-1,-1).unsqueeze(0)\n",
    "        # image_tensor_high_3ch = image_tensor_high.unsqueeze(0).expand(3,-1,-1).unsqueeze(0)\n",
    "        # # 计算ilniqe指标\n",
    "        # ilniqe_score = ilniqe_metric(image_tensor.to(device)).item()\n",
    "        # ilniqe_score_low = ilniqe_metric(image_tensor_low_3ch.to(device)).item()\n",
    "        # ilniqe_score_high = ilniqe_metric(image_tensor_high_3ch.to(device)).item()\n",
    "        \n",
    "\n",
    "\n",
    "        # 计算信息熵EN\n",
    "        EN = calculate_EN(img_int)\n",
    "        EN_low = calculate_EN(img_int_low)\n",
    "        EN_high = calculate_EN(img_int_high)\n",
    "        # 计算空间频率SF\n",
    "        SF = calculate_SF(img_double)\n",
    "        SF_low = calculate_SF(img_double_low)\n",
    "        SF_high = calculate_SF(img_double_high)\n",
    "        # 计算标准差SD\n",
    "        SD = calculate_SD(img_double)\n",
    "        SD_low = calculate_SD(img_double_low)\n",
    "        SD_high = calculate_SD(img_double_high)\n",
    "        # 计算平均梯度\n",
    "        AG = calculate_AG(img_double)\n",
    "        AG_low = calculate_AG(img_double_low)\n",
    "        AG_high = calculate_AG(img_double_high)\n",
    "        # 计算信噪比SNR\n",
    "        SNR = calculate_SNR(image_ten.to(device))\n",
    "        SNR_low = calculate_SNR(image_tensor_low.to(device))\n",
    "        SNR_high = calculate_SNR(image_tensor_high.to(device))\n",
    "        # 将结果添加到列表中\n",
    "        results_list.append({\n",
    "            'idx':i,\n",
    "            'cifar100_label': label,\n",
    "            'SNR': SNR,\n",
    "            'SNR_low': SNR_low,\n",
    "            'SNR_high': SNR_high,\n",
    "            'EN': EN,\n",
    "            'EN_low': EN_low,\n",
    "            'EN_high': EN_high,\n",
    "            'SF': SF,\n",
    "            'SF_low': SF_low,\n",
    "            'SF_high': SF_high,\n",
    "            'SD': SD,\n",
    "            'SD_low': SD_low,\n",
    "            'SD_high': SD_high,\n",
    "            'AG': AG,\n",
    "            'AG_low': AG_low,\n",
    "            'AG_high': AG_high,\n",
    "            'BRISUE': brisque_score,\n",
    "            'BRISQUE_low': brisque_score_low,\n",
    "            'BRISQUE_high': brisque_score_high,\n",
    "            # 'ILNIQE': ilniqe_score,\n",
    "            # 'ILNIQE_low': ilniqe_score_low,\n",
    "            # 'ILNIQE_high': ilniqe_score_high,\n",
    "        })\n",
    "\n",
    "    print(\"Saving file...\")\n",
    "    # 将列表转换为DataFrame，并指定列名\n",
    "\n",
    "    results_df = pd.DataFrame(results_list)\n",
    "\n",
    "    # 保存结果到csv文件中   \n",
    "    csv_file_path = Output_fil_dir\n",
    "    results_df.to_csv(csv_file_path,index=False)\n",
    "    print(\"Done.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculating image metrics...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50000/50000 [46:58<00:00, 17.74it/s]  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving file...\n",
      "Done.\n",
      "Calculating image metrics...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10000/10000 [06:14<00:00, 26.70it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving file...\n",
      "Done.\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "    generate_img_filtered_metrics(trainset,train_metrics_dir)\n",
    "    generate_img_filtered_metrics(testset,test_metrics_dir)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
